Rework host memory allocation mechanics - #1964
Conversation
rhornung67
left a comment
There was a problem hiding this comment.
@lc-hubcast approved
|
@kennyweiss @publixsubfan @BradWhitlock @bmhan12 this is a much smaller, more focused PR that addresses much of what I have been talking about regarding the default host memory allocation scheme in Axom. If you think we should include it in the Axom release, please review. Thank you. |
There was a problem hiding this comment.
Thanks for working though this @rhornung67 (and reworking it, and reworking it).
Overall, this seems like a nice compromise given the constraints and improves our memory management, but might need to be integrated with more of our infrastructure, e.g. axom::Array, sidre::Group, ...
There are some things that we should investigate before merging
(caveat: I haven't built the code, so this is based on staring at the code w/ claude):
- we might need to also update memory_management's
reallocate(MALLOC_ALLOCATOR_ID)to move the malloc handling outside the#else.
The handling ofallocate/deallocateandreallocatewas already inconsistent in axom@develop and not introduced by your changes, but the consequences are perhaps more significant now. - With
AXOM_DEFAULT_HOST_ALLOCATOR=MALLOC,Array<double> a(n)still uses the Host allocator rather than the malloc allocator since Array's default isDynamic, (defined as Umpire's current default allocator). I'm pretty sure we'd want the defaultArray<double>to followAXOM_DEFAULT_HOST_ALLOCATOR - With
AXOM_DEFAULT_HOST_ALLOCATOR=MALLOC, sidre still defaults to umpire Host (getDefaultAllocatorID()rather thandetail::getDefaultHostAllocatorID())
| { | ||
| #ifdef AXOM_USE_UMPIRE | ||
| umpire::ResourceManager& rm = umpire::ResourceManager::getInstance(); | ||
| if(allocId == MALLOC_ALLOCATOR_ID) |
There was a problem hiding this comment.
Shouldn't setDefaultAllocator also be influenced by AXOM_DEFAULT_HOST_ALLOCATOR_USES_UMPIRE_HOST ?
It seems that if we're in an Umpire build, we can never get back to the malloc allocator as the default.
E.g., consider the following in a build with Umpire:
axom::setDefaultAllocator(Host); // starts at Malloc, ends at host
axom::setDefaultAllocator(Malloc); // still ends at HostThere was a problem hiding this comment.
This method is only setting the default allocator in the Umpire ResourceManager. As long as all Axom usage honors the Axom default (Malloc or Host) when appropriate this should be OK.
Do we need to handle the case where a user calls setDefaultAllocator() with something that is neither Malloc or Host, but is still valid on the host, such as Pinned when Umpire is enabled? As I understand it, that is the main purpose of this method.
I think only the method documentation needs to be clarified, which I did.
Does that make sense now?
There was a problem hiding this comment.
This method is only setting the default allocator in the Umpire ResourceManager. As long as all Axom usage honors the Axom default (Malloc or Host) when appropriate this should be OK.
If that's the case, perhaps it should be renamed setUmpireDefaultAllocator ?
(and be a no-op in non-umpire configs)
This is the current doxygen brief for the setDefaultAllocator
\brief Sets the default memory allocator to use.
Until now, we've been using it as Axom's default allocator
(which points to Umpire when Axom is configured against Umpire).
There was a problem hiding this comment.
I modified the doxygen brief comment to indicate that it sets the default allocator in the Umpire ResourceManager, which it did before my changes IIRC. It is a no-op in non-umpire configs.
and resolve merge conflicts
| // Device memory: fill on host, then copy to device | ||
| const auto num_bytes = n * sizeof(T); | ||
| T* src = allocate<T>(num_bytes, rm.getDefaultAllocator().getId()); | ||
| T* src = allocate<T>(n, axom::detail::getDefaultHostAllocatorID()); |
There was a problem hiding this comment.
Good catch.
It's not here but in allocate() if Umpire is enabled then we check whether the passed allocator Id is a valid Umpire allocator before the malloc check. Since we're defaulting the host allocator id to MALLOC, even for Umpire-enabled builds (unless they select UMPIRE_HOST), should we swap the order of the umpire/malloc allocation checks so we can skip checking whether the allocator Id is valid for Umpire? Or, will allocate() more likely be called with allocators that are associated with Umpire?
There was a problem hiding this comment.
Good thought. I swapped the allocation check ordering.
and resolve merge conflict
…m into task/rhornung67/host-memory-api
in Axom core and Sidre Group/View behavior
|
@kennyweiss I think I addressed all the concerns in your review comments. I have done the following:
However, I maintained the distinction between
Therefore, The goal was to preserve existing semantics and not change any user facing interfaces. The only potential confusion here is that Please take a look when you have time and let me know if this satisfies your concerns. Thank you. |
kennyweiss
left a comment
There was a problem hiding this comment.
Thanks @rhornung67
It looks like there are still some issues to work through w/ sidre and the default MALLOC allocator.
Re: Array: I think we'll want the default axom::Array to use malloc when AXOM_DEFAULT_HOST_ALLOCATOR=="MALLOC" (i.e. in the default case), but I agree that it's a behavior change and is outside the scope of this PR.
| /*! | ||
| * \brief Sets the default memory allocator to use. | ||
| * \param [in] allocId the Umpire allocator id | ||
| * \brief Sets the default memory allocator for the Umpire ResourceManager. |
| axom::copy(src, origOnHost, N * sizeof(int)); | ||
|
|
||
| int* dst = axom::reallocate(src, K, dstAllocId); | ||
| EXPECT_EQ(axom::getAllocatorIDFromPointer(dst), srcAllocId); |
| - Axom's host execution-space default allocator is now a configure-time policy. The default policy is malloc, regardless | ||
| of whether Axom is configured with Umpire enabled. Umpire builds may opt into the Umpire `HOST` resource with | ||
| `-DAXOM_DEFAULT_HOST_ALLOCATOR=UMPIRE_HOST`. Runtime per-use selection remains available through existing explicit | ||
| allocator-ID arguments. |
There was a problem hiding this comment.
Please move the release notes updates to the the unreleased section above.
There was a problem hiding this comment.
Thanks for catching that. Done.
| axom::setDefaultAllocator(allocID); | ||
|
|
||
| DataStore dsPrime; | ||
| Group* rootPrime = dsPrime.getRoot(); | ||
| const int defaultHostAllocatorID = axom::detail::getDefaultHostAllocatorID(); |
There was a problem hiding this comment.
Should there be a test/check about how allocID relates to defaultHostAllocatorID, or are they not tied to each other?
Please either add a check or a comment.
There was a problem hiding this comment.
@kennyweiss please see the comment I added. A check would be wrong in general because there is no fixed relationship between allocID and defaultHostAllocatorID in this test.
There was a problem hiding this comment.
Thanks -- that was what I assumed.
| TEST_P(UmpireTest, reallocate) | ||
| { | ||
| #if defined(AXOM_USE_GPU) && defined(UMPIRE_ENABLE_CONST) | ||
| if(allocID == axom::getUmpireResourceAllocatorID(umpire::resource::Constant)) |
There was a problem hiding this comment.
Outside your PR changes, but should this line have a comment about why we're returning early?
Presumably, we can't allocate constant memory (?)
| axom::setDefaultAllocator(allocID); | ||
|
|
||
| DataStore dsPrime; | ||
| Group* rootPrime = dsPrime.getRoot(); | ||
| const int defaultHostAllocatorID = axom::detail::getDefaultHostAllocatorID(); |
There was a problem hiding this comment.
Same comment here about adding a test/check for allocID vs defaultHostAllocatorID or a comment that they're unrelated.
| const int allocators[] = {axom::getUmpireResourceAllocatorID(umpire::resource::Host) | ||
| #ifdef AXOM_USE_GPU | ||
|
|
||
| #ifdef UMPIRE_ENABLE_PINNED | ||
| , | ||
| axom::getUmpireResourceAllocatorID(umpire::resource::Pinned) | ||
| #endif | ||
|
|
||
| #ifdef UMPIRE_ENABLE_DEVICE | ||
| , | ||
| axom::getUmpireResourceAllocatorID(umpire::resource::Device) | ||
| #endif | ||
|
|
||
| #ifdef UMPIRE_ENABLE_CONST | ||
| , | ||
| axom::getUmpireResourceAllocatorID(umpire::resource::Constant) | ||
| #endif | ||
|
|
||
| #ifdef UMPIRE_ENABLE_UM | ||
| , | ||
| axom::getUmpireResourceAllocatorID(umpire::resource::Unified) | ||
| #endif | ||
|
|
||
| #endif /* defined(AXOM_USE_GPU) */ | ||
| }; |
There was a problem hiding this comment.
I assume Malloc should not be in this list since it's outside of Umpire.
Should Dynamic be in this list?
There was a problem hiding this comment.
Yes, Malloc is not in the list because it is not related to Umpire.
Dynamic is covered by the umpire::resource::Host case. When Axom is configured with Umpire enabled, Dynamic refers to Umpire Host. I did not change this. This is the way Axom works currently.
There was a problem hiding this comment.
Can you please add the following set of tests, I think the changes in this branch might cause Umpire to throw an exception w/ the new default MALLOC allocator:
TEST(sidre_group, default_allocator_getters_on_fresh_datastore)
{
sidre::DataStore ds;
sidre::Group* root = ds.getRoot();
EXPECT_NO_THROW({ root->getDefaultArrayAllocator(); });
EXPECT_NO_THROW({ root->getDefaultTupleAllocator(); });
EXPECT_NO_THROW({ root->getDefaultAllocator(); });
}This one should be ok when guarded by AXOM_USE_UMPIRE
TEST(sidre_group, default_allocator_getters_with_umpire_id)
{
auto& rm = umpire::ResourceManager::getInstance();
const int hostId = rm.getAllocator(umpire::resource::Host).getId();
sidre::DataStore ds;
sidre::Group* root = ds.getRoot();
root->setDefaultArrayAllocator(hostId);
root->setDefaultTupleAllocator(hostId);
EXPECT_NO_THROW({ EXPECT_EQ(hostId, root->getDefaultArrayAllocator().getId()); });
EXPECT_NO_THROW({ EXPECT_EQ(hostId, root->getDefaultTupleAllocator().getId()); });
}But I think this will also throw w/ the current changes:
TEST(sidre_group, default_allocator_getter_after_exec_space_id)
{
sidre::DataStore ds;
sidre::Group* root = ds.getRoot();
root->setDefaultArrayAllocator(axom::execution_space<axom::SEQ_EXEC>::allocatorID());
EXPECT_NO_THROW({ root->getDefaultArrayAllocator(); });
}There was a problem hiding this comment.
Another test to try in sidre_view:
TEST(sidre_view, allocation_honors_axom_default_allocator)
{
auto& rm = umpire::ResourceManager::getInstance();
auto pool = rm.makeAllocator<umpire::strategy::QuickPool>(
"test_pool", rm.getAllocator(umpire::resource::Host));
const int poolId = pool.getId();
const int savedId = axom::getDefaultAllocatorID();
axom::setDefaultAllocator(poolId);
sidre::DataStore ds;
sidre::View* v = ds.getRoot()->createViewAndAllocate("v", sidre::INT_ID, 100);
EXPECT_EQ(poolId, axom::getAllocatorIDFromPointer(v->getVoidPtr()));
axom::setDefaultAllocator(savedId);
}For MALLOC, it will currently give something like:
Expected equality of these values:
poolId Which is: 3
axom::getAllocatorIDFromPointer(v->getVoidPtr()) Which is: -3
For UMPIRE_HOST, it will currently give something like:
Expected equality of these values:
poolId Which is: 3
axom::getAllocatorIDFromPointer(v->getVoidPtr()) Which is: 2
There was a problem hiding this comment.
@kennyweiss regarding your comment about adding the tests above....
These expose an API mismatch that is "semantically leaky". Specifically, when Axom is configured with AXOM_DEFAULT_HOST_ALLOCATOR=MALLOC (the default in this PR), we get the following:
root->getDefaultArrayAllocatorID() returns the Sidre default allocator ID, which is Axom’s synthetic MALLOC_ALLOCATOR_ID == -3.
root->getDefaultArrayAllocator() returns an umpire::Allocator, so it only works when that stored/default ID is an actual Umpire allocator ID.
In other words, the fresh Group default ID is -3, and Umpire cannot resolve -3 with ResourceManager::getAllocator(int). So the mismatch is that the ID getter supports Axom allocator IDs, while the Umpire-object getter assumes the ID is an Umpire allocator.
That means tests should not assume these are equivalent for MALLOC:
root->getDefaultArrayAllocatorID(); // valid: can return -3
root->getDefaultArrayAllocator(); // invalid for -3
IMO, the mismatch is due to the mixing of AXOM_MALLOC and Umpire allocator IDs in Axom core which requires different semantics for IDs that are not valid for Umpire, but valid for Axom. This was done long ago. This PR is about trying to make host allocation behavior more consistent and move us to a better default -- by changing existing APIs as little as possible
There are two options:
-
Alter your test recommendations so they check the ID getters for MALLOC defaults, and only call
getDefault*Allocator()when the default ID is an actual Umpire allocator. -
Change Sidre’s
getDefault*Allocator()behavior for MALLOC_ALLOCATOR_ID, probably by returning Umpire Host as a compatibility fallback. That avoids throws, but it is the semantic leaky-ness I referred to earlier because the returned allocator would not match the stored Sidre default allocator ID.
I would go with option 1 and change the first test, for example to this:
`TEST(sidre_group, default_allocator_getters_on_fresh_datastore)
{
sidre::DataStore ds;
sidre::Group* root = ds.getRoot();
const int defaultHostAllocatorID = axom::detail::getDefaultHostAllocatorID();
EXPECT_EQ(defaultHostAllocatorID, root->getDefaultArrayAllocatorID());
EXPECT_EQ(defaultHostAllocatorID, root->getDefaultTupleAllocatorID());
EXPECT_EQ(defaultHostAllocatorID, root->getDefaultAllocatorID());
}`
Thoughts?
Others pinged as reviewers, please feel free to chime in as well.
There was a problem hiding this comment.
@kennyweiss I added variations of the Group tests you suggested based on option 1 in my comment above.
I added the default_allocator_getter_after_exec_space_id test exactly as you suggested. It does not throw because the axom::execution_space<T>::allocatorID() methods do the right thing. In general, when a method refers to ID it will do the right thing. The problem is with methods that refer to an actual allocator (i.e., Umpire allocator) and AXOM_MALLOC doesn't map to anything.
Summary
axom::MALLOC_ALLOCATOR_IDthe default host allocator when Axom is configured with Umpire enabled and when it is not.AXOM_DEFAULT_HOST_ALLOCATORwas added as a compile-time option for users who need to make a different default host execution-space policy. Its valid values areMALLOC(default) andUMPIRE_HOST(available when Axom is configured with Umpire), which was previous default.UMPIRE_HOSTcase.setDefaultAllocator()as the interface for Umpire default allocator state has been preserved.No substantial changes to Axom public APIs were made and no static state was introduced.
This PR addresses #1816